import os

import torch

from .lib import dose_stat_torch  # noqa: F401

# 加载自定义算子库
lib_path = os.path.join(os.path.dirname(__file__), "..", "..", "build", "libdose_stat_torch.so")
if os.path.exists(lib_path):
    torch.ops.load_library(lib_path)
else:
    # 如果库不存在, 尝试从当前目录加载
    lib_path = os.path.join(os.path.dirname(__file__), "libdose_stat_torch.so")
    if os.path.exists(lib_path):
        torch.ops.load_library(lib_path)


class DoseStatTorch:
    """
    PyTorch版本的剂量统计计算类
    """

    @staticmethod
    def accumulate(voxelVolume: torch.Tensor, voxelDose: torch.Tensor, doExactAccumulation: bool = False):
        """
        计算DVH(剂量体积直方图)

        Args:
            voxelVolume: 体素体积张量 (1D, double)
            voxelDose: 体素剂量张量 (1D, double)
            doExactAccumulation: 是否使用精确计算方法

        Returns:
            tuple: (binnedDoseValues, accumulatedRelativeVolumes)
        """
        # 确保输入张量类型正确
        if voxelVolume.dtype != torch.float64:
            voxelVolume = voxelVolume.double()
        if voxelDose.dtype != torch.float64:
            voxelDose = voxelDose.double()

        result = torch.ops.dose_stat_torch.accumulate(voxelVolume, voxelDose, doExactAccumulation)

        # 结果张量的形状为 [2, n], 第一行是accumulatedRelativeVolumes, 第二行是binnedDoseValues
        accumulatedRelativeVolumes = result[0]
        binnedDoseValues = result[1]

        return binnedDoseValues, accumulatedRelativeVolumes

    @staticmethod
    def get_metric(doseAtVolume: torch.Tensor, binnedDoseValues: torch.Tensor, accumulatedRelativeVolumes: torch.Tensor) -> torch.Tensor:
        """
        获取特定体积百分比对应的剂量

        Args:
            doseAtVolume: 目标体积百分比 (标量, double)
            binnedDoseValues: 分箱剂量值 (1D, double)
            accumulatedRelativeVolumes: 累积相对体积 (1D, double)

        Returns:
            torch.Tensor: 对应的剂量值
        """
        # 确保输入张量类型正确
        if doseAtVolume.dtype != torch.float64:
            doseAtVolume = doseAtVolume.double()
        if binnedDoseValues.dtype != torch.float64:
            binnedDoseValues = binnedDoseValues.double()
        if accumulatedRelativeVolumes.dtype != torch.float64:
            accumulatedRelativeVolumes = accumulatedRelativeVolumes.double()

        return torch.ops.dose_stat_torch.get_metric(doseAtVolume, binnedDoseValues, accumulatedRelativeVolumes)

    @staticmethod
    def get_dose_stat_roi_interp(dose_type: str, binnedDoseValues: torch.Tensor, accumulatedRelativeVolumes: torch.Tensor, voxelVolume: torch.Tensor, voxelDose: torch.Tensor, doseAtVolume: float = 0.0) -> torch.Tensor:
        """
        计算ROI内的剂量统计指标

        Args:
            dose_type: 剂量类型 ('average', 'D_98', 'D_95', 'D_50', 'D_2', 'DoseAtVolume')
            binnedDoseValues: 分箱剂量值 (1D, double)
            accumulatedRelativeVolumes: 累积相对体积 (1D, double)
            voxelVolume: 体素体积张量 (1D, double)
            voxelDose: 体素剂量张量 (1D, double)
            doseAtVolume: 目标体积百分比 (仅当dose_type='DoseAtVolume'时使用)

        Returns:
            torch.Tensor: 剂量统计指标值
        """
        # 确保输入张量类型正确
        if binnedDoseValues.dtype != torch.float64:
            binnedDoseValues = binnedDoseValues.double()
        if accumulatedRelativeVolumes.dtype != torch.float64:
            accumulatedRelativeVolumes = accumulatedRelativeVolumes.double()
        if voxelVolume.dtype != torch.float64:
            voxelVolume = voxelVolume.double()
        if voxelDose.dtype != torch.float64:
            voxelDose = voxelDose.double()

        return torch.ops.dose_stat_torch.get_dose_stat_roi_interp(dose_type, binnedDoseValues, accumulatedRelativeVolumes, voxelVolume, voxelDose, doseAtVolume)


# 提供与原始C++函数类似的接口
def accumulate(voxelVolume: torch.Tensor, voxelDose: torch.Tensor, doExactAccumulation: bool = False):
    """计算DVH(剂量体积直方图)"""
    return DoseStatTorch.accumulate(voxelVolume, voxelDose, doExactAccumulation)


def get_metric(doseAtVolume: torch.Tensor, binnedDoseValues: torch.Tensor, accumulatedRelativeVolumes: torch.Tensor) -> torch.Tensor:
    """获取特定体积百分比对应的剂量"""
    return DoseStatTorch.get_metric(doseAtVolume, binnedDoseValues, accumulatedRelativeVolumes)


def get_dose_stat_roi_interp(dose_type: str, binnedDoseValues: torch.Tensor, accumulatedRelativeVolumes: torch.Tensor, voxelVolume: torch.Tensor, voxelDose: torch.Tensor, doseAtVolume: float = 0.0) -> torch.Tensor:
    """计算ROI内的剂量统计指标"""
    return DoseStatTorch.get_dose_stat_roi_interp(dose_type, binnedDoseValues, accumulatedRelativeVolumes, voxelVolume, voxelDose, doseAtVolume)


# 导出接口
__all__ = ["DoseStatTorch", "accumulate", "get_metric", "get_dose_stat_roi_interp"]
